# -*- coding: utf-8 -*-
# time: 2025/4/18 16:57
# file: llm_qwen2_hf.py
# author: hanson
import gradio as gr

import torch
import os
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread

device = "cuda" if torch.cuda.is_available() else "cpu"

model_name = r"E:\soft\model\qwen\Qwen\Qwen2___5-0___5B-Instruct"

def load_model_tokenizer():
    """
    加载预训练模型和分词器。
    """
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype="auto",
        device_map="auto"
    ).eval()

    model.generation_config.max_new_tokens = 2048

    return model, tokenizer


# 大模型回答
def chat_stream(model, tokenizer, history):
    inputs = tokenizer.apply_chat_template(
        history,
        add_generation_prompt=True,
        return_tensors='pt',
    ).to(device)

    streamer = TextIteratorStreamer(tokenizer=tokenizer, skip_prompt=True, timeout=60.0, skip_special_tokens=True)
    generation_kwargs = dict(
        input_ids=inputs,
        streamer=streamer,
    )
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    for new_text in streamer:
        yield new_text


def predict(system, query, chatbot, history):
    global generation_stopped, model, tokenizer
    generation_stopped = False

    # 添加system
    if history == []:
        history.append({'role': 'system', 'content': system})
    else:
        history[0] = {'role': 'system', 'content': system}

    chatbot.append((query, ""))
    history.append({'role': 'user', 'content': query})

    response = ""
    for new_text in chat_stream(model, tokenizer, history):
        if generation_stopped:
            break
        response += new_text
        chatbot[-1] = (query, response)
        yield chatbot

    history.append({'role': 'assistant', 'content': response})


def reset_user_input():
    return gr.update(value="")


def reset_state(chatbot, history):
    history.clear()
    chatbot.clear()
    return chatbot


def regenerate(system, chatbot, history):
    # 如果没有历史记录，则直接返回
    if not history:
        yield chatbot
        return
    assistant = history.pop(-1)
    query = history.pop(-1)
    chatbot.pop(-1)
    yield from predict(system, query["content"], chatbot, history)


def stop_generation():
    global generation_stopped
    generation_stopped = True


if __name__ == '__main__':
    global model, tokenizer
    model, tokenizer = load_model_tokenizer()

    with gr.Blocks(css=os.path.dirname(__file__) + "/qwen2.css") as demo:
        gr.Markdown("""\
<p align="center"><img src="https://qianwen-res.oss-accelerate-overseas.aliyuncs.com/logo_qwen2.png" style="height: 80px"/><p>""")
        gr.Markdown(
            """\
<center><font size=3>本WebUI基于Qwen2打造，实现聊天机器人功能。</center>""")

        # 对话历史记录
        history = gr.State([])

        system = gr.Textbox(lines=2, label='System')
        chatbot = gr.Chatbot(label='Qwen2', elem_classes="control-height")

        with gr.Row():
            query = gr.Textbox(lines=2, label='Input', elem_id="query")
            submit_btn = gr.Button("🚀 Submit (发送)", elem_id="send")
            stop_btn = gr.Button("⏹️ Stop (停止)", elem_id="stop")  # 添加停止按钮

        with gr.Row():
            empty_btn = gr.Button("🧹 Clear History (清除历史)")
            regen_btn = gr.Button("🤔️ Regenerate (重试)")

        submit_btn.click(predict, [system, query, chatbot, history], [chatbot], show_progress=True)
        submit_btn.click(reset_user_input, [], [query])
        empty_btn.click(reset_state, [chatbot, history], outputs=[chatbot], show_progress=True)
        regen_btn.click(regenerate, [system, chatbot, history], [chatbot], show_progress=True)
        stop_btn.click(stop_generation, [], [chatbot])

    # 生成器必须要queue函数
    demo.queue().launch()